from fastapi import FastAPI
from netmiko import ConnectHandler
import json
import time

app = FastAPI()

# 加载设备配置
with open('devices.json') as f:
    devices = json.load(f)['devices']

# 厂商命令映射
VENDOR_COMMANDS = {
    "huawei": "display memory threshold ",
    "cisco": "show processes cpu | exclude 0.00",
    "ruijie": "show cpu core",
    "h3c": "display cpu-usage summary ",
}

@app.post("/mcp")
async def handle_request(data: dict):
    target_vendor = data.get("params", {}).get("vendor", "all")
    
    results = {}
    for dev in devices:
        if target_vendor != "all" and dev["vendor"] != target_vendor:
            continue
        
        try:
            conn = ConnectHandler(
                ip=dev["ip"],
                device_type=dev["device_type"],
                username=dev["username"],
                password=dev["password"],
                timeout=10
            )
            output = conn.send_command(VENDOR_COMMANDS.get(dev["vendor"], "show version"))
            results[dev["ip"]] = output
            conn.disconnect()
        except Exception as e:
            results[dev["ip"]] = f"Error: {str(e)}"
    
    # 生成报告
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    report_file = f"report_{target_vendor}_{timestamp}.txt"
    with open(report_file, 'w', encoding='utf-8') as f:
        f.write(f"=== {target_vendor.upper()}巡检报告 ===\n")
        f.write(json.dumps(results, indent=2, ensure_ascii=False))
    
    return {"status": "success", "report": report_file}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)